import faiss
import numpy as np
import os
import pickle
from typing import Dict, Any, List
from rich.progress import track

class FaissDB:
    def __init__(self, index_path="faiss_index"):
        self.index_path = index_path
        os.makedirs(index_path, exist_ok=True)
        self.index = self._init_index()
        self.id_to_chunk = {}
        self._load_metadata()

    def _init_index(self):
        index_file = os.path.join(self.index_path, "index.bin")
        if os.path.exists(index_file):
            return faiss.read_index(index_file)
        return faiss.IndexFlatIP(1024)  # BGE-M3维度

    def _load_metadata(self):
        meta_file = os.path.join(self.index_path, "meta.pkl")
        if os.path.exists(meta_file):
            with open(meta_file, "rb") as f:
                self.id_to_chunk = pickle.load(f)

    def _save_metadata(self):
        with open(os.path.join(self.index_path, "meta.pkl"), "wb") as f:
            pickle.dump(self.id_to_chunk, f)

    def add_chunks(self, chunks: List[Dict[str, Any]], embeddings: List[List[float]]):
        """添加分片到索引"""
        if not embeddings:
            return

        # 转换为numpy数组并归一化
        embeddings = np.array(embeddings, dtype=np.float32)
        faiss.normalize_L2(embeddings)

        # 合并索引
        if self.index.ntotal > 0:
            old_embeddings = self.index.reconstruct_n(0, self.index.ntotal)
            combined = np.vstack([old_embeddings, embeddings])
            self.index.reset()
            self.index.add(combined)
        else:
            self.index.add(embeddings)

        # 更新元数据
        start_id = len(self.id_to_chunk)
        for i, chunk in enumerate(chunks):
            self.id_to_chunk[start_id + i] = chunk

        self._save_metadata()

    def search(self, query_embedding: List[float], top_k=5) -> List[Dict[str, Any]]:
        """相似度搜索"""
        query = np.array([query_embedding], dtype=np.float32)
        faiss.normalize_L2(query)
        distances, indices = self.index.search(query, top_k)
        
        return [
            {**self.id_to_chunk[i], "score": float(d)}
            for i, d in zip(indices[0], distances[0])
            if i in self.id_to_chunk
        ]

    def save(self):
        """保存索引到磁盘"""
        faiss.write_index(self.index, os.path.join(self.index_path, "index.bin"))
        self._save_metadata()
